import numpy as np
from scipy import signal, interpolate
import tqdm
from scipy.signal import detrend,lombscargle
from scipy.optimize import curve_fit

def norm(d,ONE_AXIS=False,EPS=1e-10,BAR=False):

    if not ONE_AXIS:
        d = d-np.mean(d)
        d = d/(np.max(np.abs(d))+EPS)

        return d
    
    shape_raw = d.shape
    nt = shape_raw[-1]
    data_2d = d.reshape([-1,nt])
    ne,_ = data_2d.shape

    iter = range(ne)
    if BAR:
        iter=tqdm.tqdm(iter)

    for ie in iter:
        trace =  data_2d[ie,:]
        trace = trace-np.mean(trace)
        trace = trace/(np.max(np.abs(trace))+EPS)
        data_2d[ie,:] = trace

    return data_2d.reshape(shape_raw)

def downsample(d, t, RATE, N_BUTTER=4, Fe_LIM=5):
    nt = len(t)
    assert d.shape[-1]==nt
    d_flatten = d.reshape([-1,nt])
    
    N, nt =d_flatten.shape
    dt = t[1]-t[0]

    t_dw = t[::RATE]
    nt_dw = len(t_dw)
    d_dw = np.zeros([N,nt_dw],dtype=d_flatten.dtype)

    f_max = 1/dt/2
    fe = 1/(dt*RATE)/2-Fe_LIM
    b,a = signal.butter(N_BUTTER, fe/f_max, btype='lowpass')
    for i in tqdm.tqdm(range(N)):
        trace = d_flatten[i,:]
        trace = signal.filtfilt(b,a,trace)
        d_dw[i,:] = trace[::RATE]
    d_dw = d_dw.reshape(list(d.shape[:-1])+[nt_dw])

    return d_dw, t_dw


def filtfilt(data:np.ndarray, dt, FS, N_CORE=None, type='bandpass', BAR=False):
    
    f_max = 1/dt/2
    if type=='bandpass':
        fs,fe = FS
        b,a = signal.butter(4, [fs/f_max,fe/f_max],"bandpass")
    elif type=='lowpass':
        fe = FS
        b,a = signal.butter(4, fe/f_max,"low")
    elif type=='highpass':
        fs = FS
        b,a = signal.butter(4, fs/f_max,"high")
    else:
        raise ValueError(f'No such type value {type}')
    shape_raw = data.shape
    nt = shape_raw[-1]
    data_2d = data.reshape([-1,nt])
    
    ne,_ = data_2d.shape

    if N_CORE is None:
        iter = range(ne)
        if BAR:
            iter=tqdm.tqdm(iter)
        for ie in iter:
            trace =  data_2d[ie,:]
            data_2d[ie,:] = signal.filtfilt(b,a,trace)
        return data_2d.reshape(shape_raw)
    # else:

def avg_mean(trace, N=24, EPS=1e-10, kind='mean'):

    nt = len(trace)
    data = trace[:int(nt//N*N)].reshape([N,nt//N])
    if kind=='mean':
        d_amp = np.mean(data, axis=1)
    elif kind=='L2_mean':
        d_amp = np.sqrt(np.mean(data**2, axis=1))
    step = (np.arange(0,N)+0.5) * int(nt//N)
    t = np.arange(0,nt)
    f = interpolate.interp1d(step, d_amp,kind='slinear',bounds_error=False, fill_value=(d_amp[0],d_amp[-1]))
    mask = f(t)
    return mask

def smooth_by_mean(trace, MEAN_L = 100):
    npts = len(trace)
    mean_filter = np.ones([MEAN_L]) 
    trace_s = np.convolve(trace, mean_filter,'full')[(MEAN_L//2-1):(npts+MEAN_L//2-1)]

    return trace_s/MEAN_L


def correct_var_by_remove(data:np.ndarray, sigma_CUT=3, ITER=10, VALID_DATA=0.9,PRINT=False, t_data=None, REMOVE_VALUES=[]):

    if len(REMOVE_VALUES)>0:
        data_clean = []
        for value in data:
            if value not in REMOVE_VALUES:
                data_clean.append(value)
        data = np.array(data_clean)

    mean_0, var_0 = data.mean(), data.var()**0.5
    N_0 = len(data)

    for i in range(ITER):
        mean_i =data.mean()
        sigma_i =data.var()**0.5
        valid_idx = np.where(np.abs(data-mean_i)<sigma_CUT*sigma_i)
        if len(valid_idx[0])/N_0<VALID_DATA:
            break

        data = data[valid_idx]

        if t_data is not  None:
            t_data = t_data[valid_idx]

    mean_e, var_e = data.mean(), data.var()**0.5
    if PRINT:
        print(f"start:N {N_0}|{mean_0}|{var_0}----iter {i}---->N {len(data)}|{mean_e}|{var_e}")
    if t_data is  None:
        return data, mean_e, var_e
    else:
        return data, t_data

def gaussian_fit(X,Y):

    Y = Y/Y.sum()
    def gauss(x: np.ndarray, a: float, mu: float, sigma: float, b: float) -> np.ndarray:
        return (
            a/sigma/np.sqrt(2*np.pi)
        )*np.exp(
            -0.5 * ((x-mu)/sigma)**2
        ) + b

    mu0 = X[Y.argmax()]
    sigma0 = X.var()**0.5
    b0 = Y.mean()
    a0 = (Y.max() - b0) * sigma0 * np.sqrt(2*np.pi)
    p0 = a0, mu0, sigma0, b0

    popt, _ = curve_fit(
        f=gauss, xdata=X, ydata=Y, p0=p0
    )


    return gauss(X, *popt), popt

def linear_fit(X,Y):
    
    def linear(x, a, b):
        return a*x+b
    
    a0 = (Y.max()-Y.min())/ (X.max()-X.min())
    b0 = 0

    p0=a0, b0
    popt, _ = curve_fit(
        f=linear, xdata=X, ydata=Y, p0=p0
    )

    return linear(X, *popt), popt

def polynomial_fit(X,Y,order=2):

    c1 = np.polynomial.polynomial.polyfit(X,Y, order)
    func_fit = np.poly1d(c1[::-1])

    return func_fit(X),c1[::-1]


def detrend_unregular(data:np.ndarray,t:np.ndarray):

    data_linear, popt = linear_fit(t,data)

    return data - data_linear

def detrend_porder(data:np.ndarray,t:np.ndarray, ORDER=3):

    data_linear,_ = polynomial_fit(t,data,ORDER)
    
    return data-data_linear

def my_vel_regress(x,t, clean=True):
    '''
    自定义的线性回归函数,用于获取速度
    x: 待拟合的函数
    t: 函数对应的时间
    clean: 是否进行异常值检测和去除

    return v, t0, R

    '''
    from scipy.stats import linregress
    nx = len(x)
    slope_init, intercept_init, r_init, _, _ = linregress(t,x)
    if not clean:
        return slope_init, intercept_init, r_init
    else:

        residuals = x - (slope_init * t + intercept_init)
        mad = np.median(np.abs(residuals - np.median(residuals)))
        is_outlier = np.abs(residuals) > 3 * mad
        t_clean, x_clean = t[~is_outlier], x[~is_outlier]
        # 重新拟合
        if t_clean.size <= nx/2:
            return slope_init, intercept_init, r_init
        
        else:
            slope, intercept, r_value, p_value, std_err = linregress(t_clean, x_clean)
            return slope, intercept, r_value

# 时差测量函数
def measure_shift(trace, trace_ref, t_valid, T_window, FS=[2,4], MWCS=False,dt_interp=0.0001, PRINT=False):

    t0,t1 = T_window
    
    from obspy.signal.invsim import cosine_taper
    from scipy.signal import detrend
    fs,fe = FS
    f_max = 1/dt_interp/2
    t_interp = np.arange(t0,t1,dt_interp)
    nt = len(t_interp)
    b,a = signal.butter(1, [fs/f_max, fe/f_max],'bandpass')
    taper_mask = cosine_taper(nt, 0.85)

    f = interpolate.interp1d(t_valid,trace, bounds_error=False, fill_value=0, kind='slinear')
    trace_interp = f(t_interp)
    trace_interp = detrend(trace_interp, type='linear')
    trace_interp = taper_mask*trace_interp
    trace_interp = signal.filtfilt(b,a,trace_interp)

    f = interpolate.interp1d(t_valid,trace_ref, bounds_error=False, fill_value=0, kind='slinear')
    trace_ref_interp = f(t_interp)
    trace_ref_interp = detrend(trace_ref_interp, type='linear')
    trace_ref_interp = taper_mask*trace_ref_interp
    trace_ref_interp = signal.filtfilt(b,a,trace_ref_interp)

    coor = signal.correlate(trace_interp,trace_ref_interp)
    coor = coor/(np.sqrt( np.sum(trace_interp**2    ))+1e-10)
    coor = coor/(np.sqrt( np.sum(trace_ref_interp**2))+1e-10)
    # print(coor)
    # plot(a,linewidth=0.1)
    diff_t = (np.argmax(coor)*dt_interp-(len(t_interp)-1)*dt_interp)*1000
    if PRINT:
        print(f'time shift: delta_t:{diff_t:.2f}, corr_factor{coor.max():.4f}')
    return diff_t


def measure_shift_W(trace, trace_ref, t_valid, T_window, FS=[2,4], MWCS=False,dt_interp=0.0001, PRINT=False, W=0.1):
    
    t0,t1 = T_window
    
    from obspy.signal.invsim import cosine_taper
    from scipy.signal import detrend
    fs,fe = FS
    f_max = 1/dt_interp/2
    t_interp = np.arange(t0,t1,dt_interp)
    nt = len(t_interp)
    b,a = signal.butter(1, [fs/f_max, fe/f_max],'bandpass')
    taper_mask = cosine_taper(nt, 0.85)

    f = interpolate.interp1d(t_valid,trace, bounds_error=False, fill_value=0, kind='cubic')
    trace_interp = f(t_interp)
    trace_interp = detrend(trace_interp, type='linear')
    trace_interp = taper_mask*trace_interp
    trace_interp = signal.filtfilt(b,a,trace_interp)

    f = interpolate.interp1d(t_valid,trace_ref, bounds_error=False, fill_value=0, kind='cubic')
    trace_ref_interp = f(t_interp)
    trace_ref_interp = detrend(trace_ref_interp, type='linear')
    trace_ref_interp = taper_mask*trace_ref_interp
    trace_ref_interp = signal.filtfilt(b,a,trace_ref_interp)

    coor = signal.correlate(trace_interp,trace_ref_interp)
    coor = coor/(np.sqrt( np.sum(trace_interp**2    ))+1e-10)
    coor = coor/(np.sqrt( np.sum(trace_ref_interp**2))+1e-10)
    # print(coor)
#     diff_t = (np.argmax(coor)*dt_interp-(len(t_interp)-1)*dt_interp)*1000
    
    WP = int(W/dt_interp/2)
    coor = coor[len(t_interp)-1- WP:len(t_interp)-1+ WP]
    diff_t = (np.argmax(coor)*dt_interp-WP*dt_interp)*1000
    if PRINT:
        print(f'time shift: delta_t:{diff_t:.2f}, corr_factor{coor.max():.4f}')
    return diff_t, coor


def measure_shift_fft(trace, trace_ref, t_valid, T_window:list, FS=[2,4], P_taper=0.3, FIT=False):
    '''
    从频率域测量两道数据的时差
    trace: 待测量的波形
    trace_ref: 参考波形
    t_valid: 待测量的波形对应的时间
    T_window: 待测量的波形对应的时间范围
    FS: 待测量的波形对应的频率范围
    P_taper: 窗函数的宽度
    FIT: 是否通过拟合来计算dt

    return: diff_t,freq_valid, diff_phi
    '''

    from obspy.signal.invsim import cosine_taper
    from scipy.signal import detrend

    t0,t1 = T_window
    tp0,tp1 = np.argmin(np.abs(t_valid-t0)),np.argmin(np.abs(t_valid-t1))
    dt = t_valid[1]-t_valid[0]
    t_interp = t_valid[tp0:tp1]
    nt = len(t_interp)

    fs,fe = FS
    freq = np.fft.fftfreq(nt, d=dt)[0:int(nt//2)]
    # print(fs,fe,freq,nt,t0,t1,tp0,tp1)
    fsp, fep = np.argmin(np.abs(freq-fs)),np.argmin(np.abs(freq-fe))
    freq_valid = freq[fsp:fep]
    nf =  len(freq_valid)

    taper_mask = cosine_taper(nt, P_taper)

    trace_interp = trace[tp0:tp1]
    trace_interp = detrend(trace_interp, type='linear')
    trace_interp = taper_mask*trace_interp

    trace_ref_interp = trace_ref[tp0:tp1]
    trace_ref_interp = detrend(trace_ref_interp, type='linear')
    trace_ref_interp = taper_mask*trace_ref_interp

    spec = np.fft.fft(trace_interp)[fsp:fep]
    spec_ref = np.fft.fft(trace_ref_interp)[fsp:fep]

    
    diff_phi = np.zeros([nf])
    for i in range(nf):
        diff_phi_i = np.angle(spec[i])-np.angle(spec_ref[i])
        diff_phi_i = (diff_phi_i+np.pi)%(np.pi*2)-np.pi

        diff_phi[i] = diff_phi_i
    # diff_phi = np.unwrap(diff_phi)
    # diff_phi = diff_phi-int(diff_phi.mean()/(2*np.pi))*2*np.pi
    
    if not FIT:
        ref_A = np.abs(spec_ref)
        ref_A[np.where(ref_A<ref_A.max()/5)]=0
        weight = ref_A/(np.sum(ref_A)+1e-6)
        diff_t = -diff_phi/(2*np.pi*freq_valid)*1000
        diff_t = diff_t*weight
        diff_t = diff_t.sum()
    else:
        ref_A = np.abs(spec_ref)
        ref_A[np.where(ref_A<ref_A.max()/5)]=0
        weight = ref_A/(np.sum(ref_A)+1e-6)
        p=np.polyfit( 2*np.pi*freq_valid,-diff_phi,deg=1, w = weight)
        diff_t = p[0]*1000

    return diff_t,freq_valid, diff_phi

def measure_shift_fft_full(trace, trace_ref, t_valid, T_window:list, FS=[2,4], L_Taper=0.5,WIN_TYPE='hann', FIT=False):
    '''
    从频率域测量两道数据的时差
    
    该函数使用频域方法计算两个信号之间的时差。通过对信号进行FFT变换，在频率域中计算相位差，
    然后利用相位差与时间差的关系计算出两个信号之间的时间偏移。
    
    参数:
    trace: 待测量的波形数据 (1D array)
    trace_ref: 参考波形数据 (1D array)
    t_valid: 波形数据对应的时间轴 (1D array)
    T_window: 测量窗口的时间范围 [t_start, t_end] (list)
    FS: 频率范围 [f_min, f_max]，单位Hz (list, default=[2,4])
    L_Taper: taper窗口长度，向内拓展，单位秒 (float, default=0.5)
    WIN_TYPE: taper窗口类型 (str, default='hann')
    FIT: 是否使用线性拟合计算时差 (bool, default=False)
    
    返回:
    diff_t: 计算得到的时间差，单位毫秒 (float)
    freq_valid: 有效频率点数组 (1D array)
    diff_phi: 对应频率点的相位差 (1D array)
    trace_interp: 窗口化的待测量波形数据 (1D array)
    
    注意:
    - 当T_window超出t_valid范围时，会自动调整到有效范围内,窗口长度不变
    - 当T_window无效时(开始时间>=结束时间)，返回(0.0, [], [])
    '''

    from obspy.signal.invsim import cosine_taper
    from scipy.signal import detrend

    t0,t1 = T_window
    L_WIN = t1-t0
    t_min, t_max = t_valid[0], t_valid[-1]
    if t0 < t_min:
        t0 = t_min
        t1 = t_min+L_WIN
    if t1 > t_max:
        t1 = t_max
        t0 =t1-L_WIN
    if t0 >= t1:
        return 0.0, np.array([]), np.array([])
    if L_Taper*2>=L_WIN:
        raise ValueError('L_Taper*2>L_WIN','Taper放置在窗口内部,不能比窗口小，要调小L_Taper')

    tp0,tp1 = np.argmin(np.abs(t_valid-t0)),np.argmin(np.abs(t_valid-t1))
    dt = t_valid[1]-t_valid[0]
    nt = len(t_valid)

    fs,fe = FS
    freq = np.fft.fftfreq(nt, d=dt)[0:int(nt//2)]
    # print(fs,fe,freq,nt,t0,t1,tp0,tp1)
    fsp, fep = np.argmin(np.abs(freq-fs)),np.argmin(np.abs(freq-fe))
    freq_valid = freq[fsp:fep]
    nf =  len(freq_valid)

    # taper
    taper_mask = np.zeros_like(t_valid)
    taper_length = int(L_Taper/dt)
    taper_window = signal.get_window(WIN_TYPE, taper_length*2)
    # 前向taper
    taper_mask[tp0:tp0+taper_length] = taper_window[:taper_length]
    taper_mask[tp0+taper_length:tp1-taper_length]=1
    # 反向taper (避免两端突变)
    taper_mask[tp1-taper_length:tp1] = taper_window[taper_length:]


    trace_interp = trace.copy()
    trace_interp = detrend(trace_interp, type='linear')
    trace_interp = taper_mask*trace_interp

    trace_ref_interp = trace_ref.copy()
    trace_ref_interp = detrend(trace_ref_interp, type='linear')
    trace_ref_interp = taper_mask*trace_ref_interp

    spec = np.fft.fft(trace_interp)[fsp:fep]
    spec_ref = np.fft.fft(trace_ref_interp)[fsp:fep]

    
    diff_phi = np.zeros([nf])
    for i in range(nf):
        diff_phi_i = np.angle(spec[i])-np.angle(spec_ref[i])
        diff_phi_i = (diff_phi_i+np.pi)%(np.pi*2)-np.pi

        diff_phi[i] = diff_phi_i
    # diff_phi = np.unwrap(diff_phi)
    # diff_phi = diff_phi-int(diff_phi.mean()/(2*np.pi))*2*np.pi
    
    if not FIT:
        ref_A = np.abs(spec_ref)
        ref_A[np.where(ref_A<ref_A.max()/5)]=0
        weight = ref_A/(np.sum(ref_A)+1e-6)
        diff_t = -diff_phi/(2*np.pi*freq_valid)*1000
        diff_t = diff_t*weight
        diff_t = diff_t.sum()
    else:
        ref_A = np.abs(spec_ref)
        ref_A[np.where(ref_A<ref_A.max()/5)]=0
        weight = ref_A/(np.sum(ref_A)+1e-6)
        p=np.polyfit( 2*np.pi*freq_valid,-diff_phi,deg=1, w = weight)
        diff_t = p[0]*1000

    return diff_t,freq_valid, diff_phi,trace_interp


def remove_point_skip(x,peroid=100):
    x_new = np.zeros_like(x)

    nx = len(x_new)
    x_new[0] = x[0]
    for i in range(nx-1):
        if abs(x[i+1]-x_new[i])>peroid*0.9:
            # search = x[i+1]+np.arange(-3,3)*peroid
            # idx = np.argmin(np.abs(search-x_new[i]))
            # x_new[i+1] = search[idx]

            x_new[i+1] = x_new[i]
        else:
            x_new[i+1] = x[i+1]
    
    return x_new

def analyze_frequency_components(t, x, freqs, **args):
    """
    分析离散时间序列x = f(t)的频率成分
    
    参数:
    delta_t_shift: 时间序列数据
    te_bjt: 对应的时间戳 (BJT时间)
    freqs: 频率采样数组
    **args: 其他参数，传递给scipy.signal.lombscargle()
    
    返回:
    f: 频率数组
    Pxx_complex: 功率谱复数
    A: 幅度
    phi: 相位
    """    
    # 计算角频率
    angular_freqs = 2 * np.pi * freqs
    
    # 使用Lomb-Scargle方法计算周期图
    Pxx_complex = lombscargle(t, x.copy(), angular_freqs, normalize='amplitude',*args)
    
    # 找到最大功率对应的频率
    A = np.abs(Pxx_complex)
    phi = np.angle(Pxx_complex)
    
    return freqs, Pxx_complex, A, phi

